from argparse import ArgumentParser
import os

import torch
from scipy.io import wavfile
import numpy as np
from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.callbacks import ModelCheckpoint

from speaker_trainer import Model
from pytorch_lightning.callbacks import Callback

from speaker_trainer.backend import compute_eer
from tqdm import tqdm

class Adversarial_Attack_Helper(object):
    def __init__(self, model, alpha=3.0, num_iters=5, epsilon=15, std=120, score_save_path="log.txt", adv_save_dir="data/adv_data/", num_voting=50, device="cuda"):
        print("init Adversarial_Attack_Helper with epsilon {}".format(epsilon))
        print("init Adversarial_Attack_Helper with alpha {}".format(alpha))
        print("init Adversarial_Attack_Helper with num_voting {}".format(num_voting))
        print("Voting std is: {}".format(std))
        self.model = model
        self.alpha = alpha
        self.num_iters = num_iters
        self.epsilon = epsilon
        self.std = std
        self.adv_save_dir = adv_save_dir
        self.score_save_path = score_save_path

        if not os.path.exists(os.path.join(adv_save_dir, "wav")):
            os.makedirs(os.path.join(adv_save_dir, "wav"))

        self.trials = self.model.trials
        self.adv_trials_path = os.path.join(adv_save_dir, "adv_trials.lst")
        self.device = device
        self.num_voting = num_voting
        self.model.eval()
        if self.device == "cuda":
            self.model.cuda()

    def evaluate(self, trials=None):
        if trials is None:
            trials = self.trials
        with torch.no_grad():
            eer, th = self.model.cosine_evaluate(trials) 

    def attack(self):
        # adversarial attack example generation
        if os.path.exists(self.adv_trials_path):
            os.remove(self.adv_trials_path)
        adv_trials_file = open(self.adv_trials_path, "w")
        labels = []
        scores = []
        f = open(self.score_save_path, "w")
        for idx, item in enumerate(tqdm(self.trials)):
            label, enroll_path, adv_test_path, score = self.pgd_adversarial_attack_step(idx, item)
            adv_trials_file.write("{} {} {}\n".format(label, enroll_path, adv_test_path))
            labels.append(int(label))
            f.write(str(label) + " ")
            for s in score[0]:
                f.write(str(s) + " ")
            f.write("\n")
        eer, th = compute_eer(labels, scores)
        print("EER: {:.3f} %".format(eer*100))
        return self.adv_trials_path

    def pgd_adversarial_attack_step(self, idx, item):
        label = item[0]
        enroll_path = item[1]
        test_path = item[2]

        # load data
        samplerate, enroll_wav = self.load_wav(enroll_path)
        samplerate, test_wav = self.load_wav(test_path)
        max_delta = torch.zeros_like(test_wav).cuda()

        # init best_score and alpha
        label = int(label)
        if label == 1:
            alpha = self.alpha*(-1.0)
        else:
            alpha = self.alpha*(1.0)

        enroll_embedding = self.model.extract_speaker_embedding(enroll_wav)

        delta = torch.zeros_like(test_wav, requires_grad=True).cuda()
        org_wav = test_wav.clone()
        test_wav = test_wav.unsqueeze(0).repeat(self.num_voting, 1)
        test_wav = test_wav + torch.randn_like(test_wav) * self.std
        test_wav = torch.cat((test_wav, org_wav.unsqueeze(0)), axis=0)

        test_wav = test_wav + delta
        for t in range(self.num_iters):
            # extract test speaker embedding
            test_embedding = self.model.extract_speaker_embedding(test_wav + delta).squeeze(0)
            # cosine score
            score = enroll_embedding.mm(test_embedding.T)
            denom = torch.norm(enroll_embedding) * torch.norm(test_embedding, dim=1)
            score = score/denom
            score_loss = torch.mean(score)

            # compute grad and update delta
            score_loss.backward(retain_graph=True)
            delta.data = (delta + alpha*delta.grad.detach().sign()).clamp(-1*self.epsilon, self.epsilon)
            delta.grad.zero_()

        test_embedding = self.model.extract_speaker_embedding(test_wav+delta).squeeze(0)
        final_score = enroll_embedding.mm(test_embedding.T)
        denom = torch.norm(enroll_embedding) * torch.norm(test_embedding)
        final_score = final_score/denom
        max_delta = delta.data
        # Get Adversarial Attack wav

        adv_wav = org_wav + max_delta
        adv_wav = adv_wav.cpu().detach().numpy()
        final_score = final_score.cpu().detach().numpy()

        # save attack test wav
        idx = '%08d' % idx
        adv_test_path = os.path.join(self.adv_save_dir, "wav", idx+".wav")
        wavfile.write(adv_test_path, samplerate, adv_wav.astype(np.int16))
        return label, enroll_path, adv_test_path, final_score

    def load_wav(self, path):
        sample_rate, audio = wavfile.read(path)
        audio = torch.FloatTensor(audio)
        if self.device == "cuda":
            audio = audio.cuda()
        return sample_rate, audio

if __name__ == "__main__":
    # args
    parser = ArgumentParser()
    parser.add_argument('--alpha', help='', type=float, default=3.0)
    parser.add_argument('--num_iters', help='', type=int, default=5)
    parser.add_argument('--std', help='', type=int, default=300)
    parser.add_argument('--num_voting', help='', type=int, default=50)
    parser.add_argument('--epsilon', help='', type=int, default=15)
    parser.add_argument('--adv_save_dir', help='', type=str, default="data/test/")
    parser.add_argument('--device', help='', type=str, default="cuda")
    parser = Trainer.add_argparse_args(parser)
    parser = Model.add_model_specific_args(parser)
    args = parser.parse_args()

    args.alpha = float(args.epsilon/args.num_iters)

    model = Model(**vars(args))

    # 0. pop loss Function parameter
    state_dict = torch.load(args.checkpoint_path, map_location="cpu")["state_dict"]
    loss_weights = []
    for key, value in state_dict.items():
        if "loss" in key:
            loss_weights.append(key)
    for item in loss_weights:
        state_dict.pop(item)

    # load speaker encoder state dict and init the attack helper
    model.load_state_dict(state_dict, strict=False)
    print("initial parameter from pretrain model {}".format(args.checkpoint_path))
    helper = Adversarial_Attack_Helper(model, args.alpha, args.num_iters, args.epsilon,
            args.std, args.score_save_path, args.adv_save_dir, args.num_voting, args.device)

    print("evaluate in raw data")
    #helper.evaluate()

    print("attacking ")
    helper.attack()

